from scipy.io import loadmat
import sys

# sys.path.insert(0, "../../python")
sys.path.insert(0, '/home/xiaomin/wxm/mxnet/python')
import time
import random
import mxnet as mx
import numpy as np
import cv2 as cv
import logging
import os
import math
from datetime import datetime
from PIL import Image


def get_data_mat(folder):
    root_dir = '/home/xiaomin/wxm/Data/KaggleSeizure2016/mats'
    train_x = None
    train_y = None
    test_x = None
    nb_mat = 0
    nb_train = 0
    nb_test = 0
    for mat in os.listdir(root_dir + '/' + folder):
        begin = time.time()
        nb_mat += 1
        print mat
        print nb_mat
        if 'train' in mat:
            arr = loadmat(root_dir + '/' + folder + '/' + mat)
            x = np.expand_dims(arr['x'], 0)
            if train_x is None:
                train_x = x
            elif train_x is not None:
                train_x = np.concatenate([train_x, x], 0)

            y = np.expand_dims(arr['y'], 0)
            if train_y is None:
                train_y = y
            elif train_y is not None:
                train_y = np.concatenate([train_y, y], 0)

        elif 'test' in mat:
            arr = loadmat(root_dir + '/' + folder + '/' + mat)
            x = np.expand_dims(arr['x'], 0)
            if test_x is None:
                test_x = x
            elif test_x is not None:
                test_x = np.concatenate([test_x, x], 0)

        end = time.time()
        print 'time per mat: ' + str(end - begin)
        print '*' * 50

    print 'train_x.shape: ' + str(train_x.shape)
    print 'train_y.shape: ' + str(train_y.shape)
    print 'test_x.shape: ' + str(test_x.shape)
    return train_x, train_y, test_x


def load_params(prefix, epoch):
    """Load pre-trained models

    Args:
        prefix (str): model prefix
        epoch (int): epoch

    Returns:
        arg_params (dict): weight and bias
        aux_params (dict): other auxiliary parameters

    """
    print '%s-%04d.params' % (prefix, epoch)
    save_dict = mx.nd.load('%s-%04d.params' % (prefix, epoch))
    arg_params = {}
    aux_params = {}
    for k, v in save_dict.items():
        tp, name = k.split(':', 1)
        if tp == 'arg':
            arg_params[name] = v
            # print arg_params[name]
        if tp == 'aux':
            aux_params[name] = v
    return arg_params, aux_params


def load_params_copy(prefix, epoach, ctx):
    save_dict = mx.nd.load('%s-%04d.params' % (prefix, epoach))
    arg_params = {}
    aux_params = {}
    for k, v in save_dict.items():
        tp, name = k.split(':', 1)
        if tp == 'arg':
            arg_params[name] = mx.nd.zeros(v.shape, ctx)
            v.copyto(arg_params[name])
        if tp == 'aux':
            aux_params[name] = mx.nd.zeros(v.shape, ctx)
            v.copyto(aux_params[name])
    return arg_params, aux_params


class Speedometer(object):
    def __init__(self, batch_size, frequent=50):
        self.batch_size = batch_size
        self.frequent = frequent
        self.tic = time.time()
        self.last_count = 0

    def __call__(self, param):
        if param.nbatch % self.frequent == 0:
            speed = self.frequent * self.batch_size / (time.time() - self.tic)
            logging.info('Epoch[%d] Batch [%d]\tSpeed: %.2f samples/sec',
                         param.epoch, param.nbatch, speed)
            param.eval_metric.print_log()
            self.tic = time.time()


def get_callback(period):
    """Callback function in training

    Args:
        period (int): Every certain period, call evaluation function

    Returns:
        func: callback function
    """
    def _callback(param):
        if param.nbatch % period == 0:
            logging.info('Epoch[%d] Batch[%d]', param.epoch, param.nbatch)
            param.eval_metric.print_log()
    return _callback


def save_log(output_dir):
    """Save log in training process

    Args:
        output_dir (str): folder to save log

    """
    fmt = '%(asctime)s %(message)s'
    date_fmt = '%m-%d %H:%M:%S'
    logging.basicConfig(level=logging.INFO,
                        format=fmt,
                        datefmt=date_fmt,
                        filename=os.path.join(output_dir,
                                              'event_' + datetime.now().strftime('%Y_%m_%d_%H:%M:%S') + '.log'),
                        filemode='w')
    console = logging.StreamHandler()
    console.setLevel(logging.INFO)
    formatter = logging.Formatter(fmt=fmt, datefmt=date_fmt)
    console.setFormatter(formatter)
    logging.getLogger('').addHandler(console)


